import torch.nn as nn


class conv_block(nn.Module):
    expansion = 4

    def __init__(self, in_channels, out_channels, s=2, kernel=3):
        super(conv_block, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, stride=s, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels, kernel, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels * self.expansion, 1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion),
        )
        self.downsample = nn.Sequential(
            nn.Conv2d(in_channels, out_channels * self.expansion, 1, stride=s, padding=0, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion)
        )
        self.relu = nn.ReLU(True)

    def forward(self, x):
        x_downsample = self.downsample(x)
        x = self.layer(x)
        x = x + x_downsample
        x = self.relu(x)
        return x


class indentity_block(nn.Module):
    expansion = 4

    def __init__(self, in_channels, out_channels, kernel=3):
        super(indentity_block, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels, kernel, stride=1, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(True),
            nn.Conv2d(out_channels, out_channels * self.expansion, 1, stride=1, padding=0, bias=False),
            nn.BatchNorm2d(out_channels * self.expansion),
        )
        self.relu = nn.ReLU(True)

    def forward(self, x):
        x_downsample = x
        x = self.layer(x)
        x = x + x_downsample
        x = self.relu(x)
        return x


class Resnet50(nn.Module):
    def __init__(self, num_classes):
        super(Resnet50, self).__init__()
        self.layer0 = nn.Sequential(
            nn.Conv2d(3, 64, 7, stride=2, padding=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2, padding=1),
        )
        self.layer1 = nn.Sequential(
            conv_block(64, 64, s=1),
            indentity_block(256, 64),
            indentity_block(256, 64),
        )
        self.layer2 = nn.Sequential(
            conv_block(256, 128),
            indentity_block(512, 128),
            indentity_block(512, 128),
            indentity_block(512, 128),
        )
        self.layer3 = nn.Sequential(
            conv_block(512, 256),
            indentity_block(1024, 256),
            indentity_block(1024, 256),
            indentity_block(1024, 256),
            indentity_block(1024, 256),
            indentity_block(1024, 256),
        )
        self.layer4 = nn.Sequential(
            conv_block(1024, 512),
            indentity_block(2048, 512),
            indentity_block(2048, 512),
        )
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(2048, num_classes)

    def forward(self, x):
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x
